import sys
import os
import pygame
from copy import deepcopy

# -----------
# 配置
# -----------

WIDTH, HEIGHT = 720, 760  # 小窗口版本
BOARD_COLS, BOARD_ROWS = 9, 10
CELL_SIZE = 72
MARGIN_LEFT = 40
MARGIN_TOP = 60

BG_COLOR = (236, 235, 231)
GRID_COLOR = (0, 0, 0)
RED = (200, 60, 60)
BLACK = (20, 20, 20)
SELECT_COLOR = (0, 120, 255, 120)
HINT_COLOR = (0, 180, 0, 120)

PIECE_RADIUS = 28
FONT_PATH = "./font/JyunsaiKaai-Bold.ttf"
FONT_SIZE = 28

# 中文棋子映射
CONV_MAP = {
    'K': '将',  # 将/帅统一显示为汉字“将”
    'A': '士',
    'B': '象',
    'N': '马',
    'R': '车',
    'C': '炮',
    'P': '兵',
}

FILES = 'abcdefgh'
RANKS = [str(i) for i in range(1, 11)]


# -----------
# 数据结构
# -----------

class Piece:
    def __init__(self, name, color):
        self.name = name  # 'K','A','B','N','R','C','P'
        self.color = color  # 'r' 红, 'b' 黑
        self.text = CONV_MAP.get(self.name, self.name)

    def __repr__(self):
        return f"{self.color}:{self.name}"


class Board:
    def __init__(self):
        self.grid = [[None for _ in range(9)] for _ in range(10)]
        self.init_pieces()

    def init_pieces(self):
        self.grid[0][4] = Piece('K', 'b')  # 黑方将
        self.grid[9][4] = Piece('K', 'r')  # 红方将

        self.grid[0][3] = Piece('A', 'b')
        self.grid[0][5] = Piece('A', 'b')
        self.grid[9][3] = Piece('A', 'r')
        self.grid[9][5] = Piece('A', 'r')

        self.grid[0][1] = Piece('N', 'b')
        self.grid[0][7] = Piece('N', 'b')
        self.grid[9][1] = Piece('N', 'r')
        self.grid[9][7] = Piece('N', 'r')

        self.grid[2][1] = Piece('C', 'b')
        self.grid[2][7] = Piece('C', 'b')
        self.grid[7][1] = Piece('C', 'r')
        self.grid[7][7] = Piece('C', 'r')

        self.grid[0][0] = Piece('R', 'b')
        self.grid[0][8] = Piece('R', 'b')
        self.grid[9][0] = Piece('R', 'r')
        self.grid[9][8] = Piece('R', 'r')

        self.grid[0][2] = Piece('B', 'b')
        self.grid[0][6] = Piece('B', 'b')
        self.grid[9][2] = Piece('B', 'r')
        self.grid[9][6] = Piece('B', 'r')

        for c in range(0, 9, 2):
            self.grid[3][c] = Piece('P', 'b')
            self.grid[6][c] = Piece('P', 'r')

    def clone(self):
        nb = Board()
        nb.grid = deepcopy(self.grid)
        return nb

    def in_bounds(self, r, c):
        return 0 <= r < 10 and 0 <= c < 9

    def get(self, r, c):
        if self.in_bounds(r, c):
            return self.grid[r][c]
        return None

    def set(self, r, c, p):
        if self.in_bounds(r, c):
            self.grid[r][c] = p

    def move_piece(self, sr, sc, tr, tc):
        piece = self.get(sr, sc)
        self.set(tr, tc, piece)
        self.set(sr, sc, None)

    def find_king(self, color):
        for r in range(10):
            for c in range(9):
                p = self.grid[r][c]
                if p and p.color == color and p.name == 'K':
                    return (r, c)
        return None


# -----------
# 规则与判定
# -----------

def is_adjacent(sr, sc, tr, tc):
    return max(abs(tr - sr), abs(tc - sc)) == 1


def in_palace(r, c, color):
    # 简化九宫格边界：红方在下方九宫，黑方在上方九宫
    if color == 'r':
        return 7 <= r <= 9 and 3 <= c <= 5
    else:
        return 0 <= r <= 2 and 3 <= c <= 5


def can_move_piece(board, sr, sc, tr, tc, check_check=True):
    if not board.in_bounds(sr, sc) or not board.in_bounds(tr, tc):
        return False
    piece = board.get(sr, sc)
    target = board.get(tr, tc)
    if not piece:
        return False
    if target and target.color == piece.color:
        return False

    dr = tr - sr
    dc = tc - sc

    # 将/帅：不能斜着走，必须在九宫内且上下左右单步
    if piece.name == 'K':
        if not in_palace(tr, tc, piece.color):
            return False
        if (abs(dr) + abs(dc)) == 1:
            return True
        return False

    # 士：斜着移动1格，需在九宫内（保留传统士规则）
    if piece.name == 'A':
        if abs(dr) == 1 and abs(dc) == 1 and in_palace(tr, tc, piece.color):
            return True
        return False

    # 象：两格对角，需“象眼”为中点空且不得过河
    if piece.name == 'B':
        if abs(dr) == 2 and abs(dc) == 2:
            er, ec = sr + dr // 2, sc + dc // 2
            if board.get(er, ec) is None:
                # 不越河
                if piece.color == 'r' and tr >= 5:
                    return False
                if piece.color == 'b' and tr <= 4:
                    return False
                return True
        return False

    # 马：绊马脚
    if piece.name == 'N':
        if (abs(dr), abs(dc)) in ((2, 1), (1, 2)):
            if abs(dr) == 2 and abs(dc) == 1:
                leg_r, leg_c = sr + dr // 2, sc
            else:
                leg_r, leg_c = sr, sc + dc // 2
            if board.get(leg_r, leg_c) is None:
                return True
        return False

    # 车：直线，不能越过棋子
    if piece.name == 'R':
        if dr == 0 and dc != 0:
            step = 1 if dc > 0 else -1
            for cc in range(sc + step, tc, step):
                if board.get(sr, cc) is not None:
                    return False
            return True
        if dc == 0 and dr != 0:
            step = 1 if dr > 0 else -1
            for rr in range(sr + step, tr, step):
                if board.get(rr, sc) is not None:
                    return False
            return True
        return False

    # 炮：走法与车相同；吃子时需要恰好一个炮架（中间有棋子），距离不限；不吃时同样沿直线移动但需要炮架才能跨越
    if piece.name == 'C':
        if dr == 0 and dc != 0:
            step = 1 if dc > 0 else -1
            cnt = 0
            for cc in range(sc + step, tc, step):
                if board.get(sr, cc) is not None:
                    cnt += 1
            if cnt != 1:
                return False
            return True
        if dc == 0 and dr != 0:
            step = 1 if dr > 0 else -1
            cnt = 0
            for rr in range(sr + step, tr, step):
                if board.get(rr, sc) is not None:
                    cnt += 1
            if cnt != 1:
                return False
            return True
        return False

    # 兵/卒：未过河前只能直走1格，不能横向移动
    if piece.name == 'P':
        if piece.color == 'r':
            # 红方向上走
            if sc == tc:
                if dr == -1:
                    return True
            # 未越河，禁止横向
            if sr <= 4:
                return False
            # 过河后横向移动1格，且仍可直走
            if sr >= 5:
                if dr == -1 and dc == 0:
                    return True
                if dr == 0 and abs(dc) == 1:
                    return True
            return False
        else:
            # 黑方向下走
            if sc == tc:
                if dr == 1:
                    return True
            if sr >= 5:
                return False  # 已经越河，纵向仍可，但逻辑保留
            # 未越河，禁止横向
            if sr <= 4:
                return False
            if sr >= 5:
                if dr == 1 and dc == 0:
                    return True
                if dr == 0 and abs(dc) == 1:
                    return True
            return False

    return False


def is_in_check(board, color):
    king_pos = board.find_king(color)
    if not king_pos:
        return True
    kr, kc = king_pos
    opp = 'b' if color == 'r' else 'r'
    for r in range(10):
        for c in range(9):
            p = board.get(r, c)
            if p and p.color == opp:
                if can_move_piece(board, r, c, kr, kc, check_check=False):
                    return True
    return False


def is_checkmated(board, color):
    if not is_in_check(board, color):
        return False
    for r in range(10):
        for c in range(9):
            p = board.get(r, c)
            if p and p.color == color:
                for tr in range(10):
                    for tc in range(9):
                        if (r, c) == (tr, tc):
                            continue
                        if can_move_piece(board, r, c, tr, tc):
                            nb = board.clone()
                            nb.move_piece(r, c, tr, tc)
                            if not is_in_check(nb, color):
                                return False
    return True


# -----------
# 绘制与字体
# -----------

def load_font():
    if os.path.exists(FONT_PATH):
        try:
            return pygame.font.Font(FONT_PATH, FONT_SIZE)
        except Exception:
            pass
    return pygame.font.SysFont(None, FONT_SIZE)


def draw_board(screen, origin_x, origin_y, cell_size, board, selected, turn, font):
    screen.fill(BG_COLOR)

    # 网格
    for r in range(BOARD_ROWS):
        for c in range(BOARD_COLS):
            x = origin_x + c * cell_size
            y = origin_y + r * cell_size
            rect = pygame.Rect(x, y, cell_size, cell_size)
            pygame.draw.rect(screen, (210, 190, 150), rect)
            pygame.draw.rect(screen, GRID_COLOR, rect, 1)

    river_y = origin_y + 4 * cell_size
    pygame.draw.line(screen, GRID_COLOR, (origin_x, river_y),
                     (origin_x + 8 * cell_size, river_y), 2)

    if selected:
        sr, sc = selected
        x = origin_x + sc * cell_size
        y = origin_y + sr * cell_size
        sel_rect = pygame.Rect(x, y, cell_size, cell_size)
        s = pygame.Surface((cell_size, cell_size), pygame.SRCALPHA)
        s.fill(SELECT_COLOR)
        screen.blit(s, (x, y))

    if selected:
        sr, sc = selected
        for tr in range(10):
            for tc in range(9):
                if can_move_piece(board, sr, sc, tr, tc):
                    x = origin_x + tc * cell_size + cell_size // 2
                    y = origin_y + tr * cell_size + cell_size // 2
                    pygame.draw.circle(screen, HINT_COLOR, (x, y), 6)

    # 棋子绘制
    for r in range(BOARD_ROWS):
        for c in range(BOARD_COLS):
            p = board.get(r, c)
            if p:
                cx = origin_x + c * cell_size + cell_size // 2
                cy = origin_y + r * cell_size + cell_size // 2
                color = RED if p.color == 'r' else BLACK
                pygame.draw.circle(screen, color, (cx, cy), PIECE_RADIUS)
                text = font.render(p.text, True, (255, 255, 255))
                text_rect = text.get_rect(center=(cx, cy))
                screen.blit(text, text_rect)

    info_font = pygame.font.SysFont(None, 22)
    turn_text = "红方" if turn == 'r' else "黑方"
    info = info_font.render(f"当前：{turn_text}  -  先手请选棋子，选中后再点击目标格移动", True, (0, 0, 0))
    screen.blit(info, (10, HEIGHT - 28))


# -----------
# 主程序
# -----------

def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("中国象棋 - pygame 实现（严格修正版）")
    clock = pygame.time.Clock()

    font = load_font()

    origin_x = MARGIN_LEFT
    origin_y = MARGIN_TOP
    cell_size = CELL_SIZE

    board = Board()
    current = 'r'
    selected = None

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                mx, my = pygame.mouse.get_pos()
                if origin_x <= mx <= origin_x + BOARD_COLS * cell_size and \
                        origin_y <= my <= origin_y + BOARD_ROWS * cell_size:
                    c = (mx - origin_x) // cell_size
                    r = (my - origin_y) // cell_size
                    if not (0 <= r < 10 and 0 <= c < 9):
                        continue
                    piece = board.get(r, c)
                    if selected is None:
                        if piece and piece.color == current:
                            selected = (r, c)
                    else:
                        sr, sc = selected
                        tr, tc = r, c
                        if (sr, sc) == (tr, tc):
                            selected = None
                        else:
                            if can_move_piece(board, sr, sc, tr, tc):
                                board.move_piece(sr, sc, tr, tc)
                                opponent = 'b' if current == 'r' else 'r'
                                if is_in_check(board, opponent) and is_checkmated(board, opponent):
                                    print(f"{'红方' if current == 'r' else '黑方'} 将死，对局结束，重新开始。")
                                    board = Board()
                                    current = 'r'
                                    selected = None
                                    continue
                                current = opponent
                                selected = None
                            else:
                                # 无效落点，保持选中
                                pass

        screen.fill((0, 0, 0))
        draw_board(screen, origin_x, origin_y, cell_size, board, selected, current, font)

        tip_font = pygame.font.SysFont(None, 20)
        tip = tip_font.render(
            "规则要点：兵未过河前需直走，不能横向；过河后可横向移动1格；炮走法如车，吃子需炮架且沿直线，距离不限；将/帅不能斜着走，士在九宫；落点以点标示。",
            True, (0, 0, 0)
        )
        screen.blit(tip, (10, HEIGHT - 52))

        pygame.display.flip()
        clock.tick(30)

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()
